/*
 * paperback: paper backup generator suitable for long-term storage
 * Copyright (C) 2018-2022 Aleksa Sarai <cyphar@cyphar.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use crate::shamir::{
    gf::{EvaluablePolynomial, GfBarycentric, GfElem, GfElemPrimitive, GfPolynomial},
    shard::Shard,
    Error,
};

use std::mem;

use rayon::prelude::*;

/// Factory to share a secret using [Shamir Secret Sharing][sss].
///
/// [sss]: https://en.wikipedia.org/wiki/Shamir%27s_Secret_Sharing
#[derive(Clone, Debug)]
pub struct Dealer {
    polys: Vec<Box<dyn EvaluablePolynomial>>,
    secret_len: usize,
    threshold: GfElemPrimitive,
}

impl Dealer {
    /// Returns the number of *unique* `Shard`s generated by this `Dealer`
    /// required to recover the stored secret.
    #[allow(dead_code)]
    pub fn threshold(&self) -> u32 {
        self.threshold
    }

    /// Construct a new `Dealer` to shard the `secret`, requiring at least
    /// `threshold` shards to reconstruct the secret.
    pub fn new<B: AsRef<[u8]>>(threshold: u32, secret: B) -> Self {
        assert!(threshold > 0, "must at least have a threshold of one");
        let k = threshold - 1;
        let secret = secret.as_ref();
        let polys = secret
            // Generate &[u32] from &[u8], by chunking into sets of four.
            .par_chunks(mem::size_of::<GfElemPrimitive>())
            .map(GfElem::from_bytes)
            // Generate a random polynomial with the value as the constant.
            .map(|x0| {
                let mut poly = GfPolynomial::new_rand(k, &mut rand::thread_rng());
                *poly.constant_mut() = x0;
                Box::new(poly) as Box<dyn EvaluablePolynomial>
            })
            .collect::<Vec<_>>();
        Dealer {
            polys,
            threshold,
            secret_len: secret.len(),
        }
    }

    /// Get the secret value stored by the `Dealer`.
    pub fn secret(&self) -> Vec<u8> {
        let mut secret = self
            .polys
            .par_iter()
            .map(|poly| poly.constant())
            .flat_map(|x| x.to_bytes())
            .collect::<Vec<_>>();

        // Cannot call .take() on rayon::iter::FlatMap, so do it the
        // old-fashioned way instead.
        secret.drain(self.secret_len..);
        secret
    }

    /// Generate a new `Shard` for the secret.
    ///
    /// NOTE: The `x` value is calculated randomly, which means that there is a
    ///       small chance that two separate calls to `Dealer::shard` will
    ///       generate the same `Shard`. It is up to the caller to be sure that
    ///       they have enough *unique* shards to reconstruct the secret.
    // TODO: I'm not convinced the chances of collision are low enough...
    pub fn next_shard(&self) -> Shard {
        let mut g = rand::thread_rng();
        // TODO: We should probably add some limit to this.
        loop {
            match self.shard(GfElem::new_rand(&mut g)) {
                Some(shard) => return shard,
                None => continue,
            }
        }
    }

    /// Generate a `Shard` for the secret using the given `x` value.
    pub fn shard(&self, x: GfElem) -> Option<Shard> {
        if x == GfElem::ZERO {
            return None;
        }

        self.polys
            .par_iter()
            .map(|poly| {
                let y = poly.evaluate(x);
                // If we are given an x value where one of the polynomials
                // happens to produce the secret, it seems produent to not leak
                // that information. In the case of next_shard(), this will just
                // cause us to generate a different x value. This happens very
                // rarely, and we don't care about paperback-expand usecases
                // because we have never allowed these broken shards to be
                // produed and thus generating one intentionally would never be
                // needed by real users.
                if y != poly.constant() || self.threshold == 1 {
                    Some(y)
                } else {
                    None
                }
            })
            .collect::<Option<Vec<_>>>()
            .map(|ys| Shard {
                x,
                ys,
                threshold: self.threshold,
                secret_len: self.secret_len,
            })
    }

    /// Reconstruct an entire `Dealer` from a *unique* set of `Shard`s.
    ///
    /// The caller must pass exactly the correct number of shards.
    ///
    /// This operation is significantly slower than `recover_secret`, so it
    /// should only be used if it is necessary to construct additional shards
    /// with `Dealer::next_shard`.
    pub fn recover<S: AsRef<[Shard]>>(shards: S) -> Result<Self, Error> {
        let shards = shards.as_ref();
        assert!(!shards.is_empty(), "must be provided at least one shard");

        let threshold = shards[0].threshold();
        let polys_len = shards[0].ys.len();
        let secret_len = shards[0].secret_len;

        // TODO: Implement this consistency checking more nicely.
        for shard in shards {
            assert!(shard.threshold() == threshold, "shards must be consistent");
            assert!(shard.ys.len() == polys_len, "shards must be consistent");
            assert!(shard.secret_len == secret_len, "shards must be consistent");
        }

        assert!(
            shards.len() == threshold as usize,
            "must have exactly {} shards",
            threshold
        );

        let polys = (0..polys_len)
            .into_par_iter()
            .map(|i| {
                let xs = shards.iter().map(|s| s.x);
                let ys = shards.iter().map(|s| s.ys[i]);

                let points = xs.zip(ys).collect::<Vec<_>>();
                GfBarycentric::recover(threshold - 1, points.as_slice())
                    .map(|poly| Box::new(poly) as Box<dyn EvaluablePolynomial>)
            })
            .collect::<Result<Vec<_>, _>>()?;

        Ok(Self {
            polys,
            secret_len,
            threshold,
        })
    }
}

#[cfg(test)]
mod test {
    use super::*;

    use quickcheck::TestResult;

    // We use u16s and u8s here (and limit the range) because we cannot handle
    // ridiculously large quorum sizes in quickcheck tests. Some larger quorums
    // are tested in v0::test::paperback_expand_smoke.

    #[quickcheck]
    fn basic_roundtrip(n: u16, secret: Vec<u8>) -> TestResult {
        if !(1..=4096).contains(&n) {
            return TestResult::discard();
        }
        let dealer = Dealer::new(n.into(), &secret);
        TestResult::from_bool(secret == dealer.secret())
    }

    #[cfg(debug_assertions)] // not --release
    const SECRET_UPPER: u8 = 64;
    #[cfg(not(debug_assertions))] // --release
    const SECRET_UPPER: u8 = 255;

    #[quickcheck]
    fn recover_secret_fail(n: u8, secret: Vec<u8>) -> TestResult {
        // Fail to recover the secret with invalid shards.
        //
        // Note that large n values take a very long time to recover the secret.
        // This is proportional to secret.len(), which is controlled by
        // quickcheck and thus can be quite large.
        if !(2..=SECRET_UPPER).contains(&n) || secret.is_empty() {
            return TestResult::discard();
        }

        let dealer = Dealer::new(n.into(), &secret);
        let shards = (0..(n - 1))
            .map(|_| {
                let mut shard = dealer.next_shard();
                shard.threshold -= 1;
                // Ensure shard IDs are always ID_LENGTH.
                assert_eq!(shard.id().len(), Shard::ID_LENGTH);
                shard
            })
            .collect::<Vec<_>>();

        TestResult::from_bool(Dealer::recover(shards).unwrap().secret() != secret)
    }

    #[quickcheck]
    fn recover_secret_success(n: u8, secret: Vec<u8>) -> TestResult {
        // Recover just the secret.
        //
        // Note that large n values take a very long time to recover the secret.
        // This is proportional to secret.len(), which is controlled by
        // quickcheck and thus can be quite large.
        if !(1..=SECRET_UPPER).contains(&n) {
            return TestResult::discard();
        }

        let dealer = Dealer::new(n.into(), &secret);
        let shards = (0..n)
            .map(|_| {
                let shard = dealer.next_shard();
                // Ensure shard IDs are always ID_LENGTH.
                assert_eq!(shard.id().len(), Shard::ID_LENGTH);
                shard
            })
            .collect::<Vec<_>>();

        TestResult::from_bool(Dealer::recover(shards).unwrap().secret() == secret)
    }

    #[cfg(debug_assertions)] // not --release
    const RECOVER_UPPER: u8 = 32;
    #[cfg(not(debug_assertions))] // --release
    const RECOVER_UPPER: u8 = 255;

    #[quickcheck]
    fn limited_recover_fail(n: u8, secret: Vec<u8>, test_xs: Vec<GfElem>) -> TestResult {
        use std::collections::HashSet;

        // Fail to recover the dealer with invalid shards.
        //
        // Note that even moderately large n values take a longer time to fully
        // recover -- which when paired with quickcheck makes it take far too
        // long. This is proportional to secret.len() (probably O(L*n^2) with
        // big constants or something like that).
        if !(2..=RECOVER_UPPER).contains(&n) || secret.is_empty() || test_xs.contains(&GfElem::ZERO)
        {
            return TestResult::discard();
        }
        let dealer = Dealer::new(n.into(), secret);
        let shards = (0..(n - 1))
            .map(|_| {
                let mut shard = dealer.next_shard();
                shard.threshold -= 1;
                // Ensure shard IDs are always ID_LENGTH.
                assert_eq!(shard.id().len(), Shard::ID_LENGTH);
                shard
            })
            .collect::<Vec<_>>();
        let shard_xs = {
            let mut xs = HashSet::new();
            for shard in &shards {
                xs.insert(shard.x);
            }
            xs
        };
        let recovered_dealer = Dealer::recover(shards).unwrap();

        println!(
            "GF(0) => dealer: {:?}, recovered_dealer: {:?}",
            dealer.shard(GfElem::ZERO),
            recovered_dealer.shard(GfElem::ZERO)
        );

        TestResult::from_bool(
            dealer.secret() != recovered_dealer.secret()
                && test_xs
                    .iter()
                    // If it is a shard x value then it will match, otherwise it
                    // should not match (because it's the wrong polynomial).
                    .all(|&x| {
                        shard_xs.contains(&x) == (dealer.shard(x) == recovered_dealer.shard(x))
                    }),
        )
    }

    #[quickcheck]
    fn limited_recover_success(n: u8, secret: Vec<u8>, test_xs: Vec<GfElem>) -> TestResult {
        // Recover the dealer so we can construct new shards.
        //
        // Note that even moderately large n values take a longer time to fully
        // recover -- which when paired with quickcheck makes it take far too
        // long. This is proportional to secret.len() (probably O(L*n^2) with
        // big constants or something like that).
        if !(1..=RECOVER_UPPER).contains(&n) || test_xs.contains(&GfElem::ZERO) {
            return TestResult::discard();
        }
        let dealer = Dealer::new(n.into(), secret);
        let shards = (0..n)
            .map(|_| {
                let shard = dealer.next_shard();
                // Ensure shard IDs are always ID_LENGTH.
                assert_eq!(shard.id().len(), Shard::ID_LENGTH);
                shard
            })
            .collect::<Vec<_>>();
        let recovered_dealer = Dealer::recover(shards).unwrap();

        TestResult::from_bool(
            dealer.secret() == recovered_dealer.secret()
                && test_xs
                    .iter()
                    .all(|&x| dealer.shard(x) == recovered_dealer.shard(x)),
        )
    }
}
